import torch
import numpy as np
from PIL import Image
from torchvision import transforms,models

model = models.vgg16(pretrained = True)
# 构建新的全连接层
model.classifier = torch.nn.Sequential(torch.nn.Linear(25088, 100),
                                       torch.nn.ReLU(),
                                       torch.nn.Dropout(p=0.5),
                                       torch.nn.Linear(100, 2))

model.load_state_dict(torch.load('cat_dog_fc.pth'))

model.eval()

label = np.array(['cat','dog'])

# 数据预处理
transform = transforms.Compose([
    transforms.Resize(224),
    transforms.ToTensor()
])

def predict(image_path):
    # 打开图片
    img = Image.open(image_path)
    # 数据处理，再增加一个维度
    img = transform(img).unsqueeze(0)
    # 预测得到结果
    outputs = model(img)
    # 获得最大值所在位置
    _, predicted = torch.max(outputs,1)
    # 转化为类别名称
    print(label[predicted.item()])

predict('image/test/cat/cat.1066.jpg')
predict('image/test/cat/cat.1097.jpg')
predict('image/test/cat/cat.1021.jpg')
predict('image/test/cat/cat.1056.jpg')
predict('image/test/cat/cat.1034.jpg')

predict('image/test/dog/dog.1024.jpg')
predict('image/test/dog/dog.1076.jpg')
predict('image/test/dog/dog.1018.jpg')
predict('image/test/dog/dog.1054.jpg')
predict('image/test/dog/dog.1066.jpg')
